package middlewares

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/juju/ratelimit"
)

//限流的中间件

func RateLimitMiddleware(fillInterval time.Duration, cap int64) func(ctx *gin.Context) {
	bucket := ratelimit.NewBucket(fillInterval, cap) //fillInterval 是填充的速率，cap 是总容量
	return func(ctx *gin.Context) {
		//如果取不到令牌就返回响应
		if bucket.TakeAvailable(1) == 0 { //取一个令牌，如果没有足够的令牌就返回 0。
			ctx.String(http.StatusOK, "rate limit...")
			ctx.Abort()
			return
		}
		//取到令牌就放行
		ctx.Next()
	}
}
